# To impute missing values
lm_impute <- function(df, missing_var, rhs, education) {
  
  edu_df <- df |> 
    filter(uddannelse == education)
  
  # Indices of rows with missing values
  missing_rows <- which(is.na(df[[missing_var]]) & df$uddannelse == education)
  
  if (nrow(edu_df) > length(missing_rows)) {
    
    # Linear model
    temp_model <- lm(
      formula = as.formula(paste(missing_var, "~", rhs)),
      data = edu_df,
      na.action = na.exclude
    )
    
    # Using model to impute missing values
    df[[missing_var]][missing_rows] <- predict(temp_model, df[missing_rows,])
    
  }
  
  return(df)
}

# For reporting
report_results <- function(fit, education_name, t0, pre_se = "bootstrapped") {
  
  # Data frames with complete observations for pretty plotting
  if (education_name == "Folkeskolelærer") {
    complete_df <- teacher_df |> 
      filter(år <= 2019 & !is.na(ansøgninger_main)) |> 
      group_by(uddannelse) |> 
      summarise(n = n()) |> 
      filter(n == max(n))
    
  } else if (education_name == "Sygeplejerske (kvote 1)") {
    complete_df <- nurse_df_k1 |> 
      filter(!is.na(ansøgninger_main)) |> 
      group_by(uddannelse) |> 
      summarise(n = n()) |> 
      filter(n == max(n))
    
  } else if (education_name == "Sygeplejerske (kvote 2)") {
    complete_df <- nurse_df_k2 |>
      filter(!is.na(ansøgninger_main)) |> 
      group_by(uddannelse) |> 
      summarise(n = n()) |> 
      filter(n == max(n))
    
  } else {
    complete_df <- NULL
  }
  
  # First making counterfactual data frame
  ct_df <- data.frame(
    education = education_name,
    t = as.vector(fit[["time"]]),
    observed = as.vector(fit[["Y.tr"]]),
    sc = as.vector(fit[["Y.ct"]]),
    t0 = t0)
  
  # Making controls data frame
  controls_df <- cbind(ct_df, fit[["Y.dat"]]) |> 
    pivot_longer(cols = -c(education, t, observed, sc, t0),
                 names_to = "control_education",
                 values_to = "applications") |> 
    filter(control_education != education_name)
  
  # Only including complete data series in plot
  if (!is.null(complete_df)) {
    controls_df <- controls_df |> 
      filter(control_education %in% complete_df$uddannelse)
  }
  
  # Making gap data frame
  gap_df <- ct_df |> 
    mutate(estimate = observed - sc) |> 
    bind_cols(as.data.frame(fit[["est.att"]]))
  
  if (pre_se == "mspe") {
    gap_df <- gap_df |> 
      mutate(
        CI.lower = if_else(
          n.Treated == 0,
          estimate - 1.96 * sqrt(fit[["MSPE"]] / sum(as.vector(fit[["est.att"]][,"n.Treated"]) == 0)),
          CI.lower),
        CI.upper = if_else(
          n.Treated == 0,
          estimate + 1.96 * sqrt(fit[["MSPE"]] / sum(as.vector(fit[["est.att"]][,"n.Treated"]) == 0)),
          CI.upper))
  }
  
  # Making plots
  ct_plot <- ct_df |> 
    ggplot(aes(y = observed, x = t)) +
    geom_line(
      aes(linetype = "Observed"),
      color = "black",
      linewidth = 1) +
    geom_line(
      aes(y = sc, linetype = "Synthetic"),
      linewidth = 1) +
    geom_vline(aes(xintercept = t0), linetype = "dotted") +
    geom_line(
      data = controls_df,
      aes(x = t, y = applications, group = control_education, linetype = "Donor education programs"),
      inherit.aes = FALSE,
      color = "grey",
      alpha = 0.5) +
    scale_linetype_manual(
      values = c(
        "Observed" = "solid",
        "Synthetic" = "dashed",
        "Donor education programs" = "solid")) +
    scale_x_continuous(expand = c(0, 0)) +
    theme_bw() +
    theme(legend.position = "bottom") +
    labs(
      x = NULL,
      y = "Applications",
      linetype = NULL) +
    ggtitle("A. Observed and counterfactual development")
  
  gap_plot <- ggplot(gap_df, aes(x = t, y = estimate)) +
    geom_line(color = "black", linewidth = 1) +
    geom_ribbon(aes(ymin = CI.lower, ymax = CI.upper), alpha = 0.3) + 
    geom_hline(yintercept = 0) +
    geom_vline(aes(xintercept = t0), linetype = "dotted") +
    scale_x_continuous(expand = c(0, 0)) +
    theme_bw() +
    labs(x = NULL,
         y = "Estimate") +
    ggtitle("B. Effect estimates")
  
  return(list(ct_plot, gap_plot))
  
}
